from tvm_graph.graph import Graph
from tvm_graph.node import Node


class GraphPattern:
    
    @staticmethod
    def get_selu():
        test_graph = Graph()
        
    
    @staticmethod
    def get_bn_graph():
        test_graph = Graph()
        node_1 = Node("add", input_num=1, out_shape=[200])
        node_1.atomic_type = "add"
        node_2 = Node("sqrt", input_num=1, out_shape=[200])
        node_2.atomic_type = "sqrt"
        node_3 = Node("divide", input_num=1, out_shape=[200])
        node_3.atomic_type = "divide"
        node_4 = Node("multiply", input_num=1, out_shape=[200])
        node_4.atomic_type = "multiply"
        node_5 = Node("multiply-2", input_num=1, out_shape=[200])
        node_5.atomic_type = "multiply"
        node_6 = Node("add-2", input_num=1, out_shape=[200])
        node_6.atomic_type = "add"
        node_7 = Node("add-3", input_num=1, out_shape=[200])
        node_7.atomic_type = "add"

        node_8 = Node("neg", input_num=1, out_shape=[200])
        node_8.atomic_type = "negative"
        node_9 = Node("multiply-3", input_num=1, out_shape=[200])
        node_9.atomic_type = "multiply"

        node_1.add_next_node(node_2)
        node_2.add_father_node(node_1)
        node_2.add_next_node(node_3)
        node_3.add_father_node(node_2)
        node_3.add_next_node(node_4)
        node_4.add_father_node(node_3)
        node_4.add_next_node(node_5)
        node_5.add_father_node(node_4)
        node_5.add_next_node(node_6)
        node_6.add_father_node(node_5)
        node_6.add_next_node(node_7)
        node_7.add_father_node(node_6)

        node_5.add_father_node(node_8)
        node_8.add_next_node(node_5)
        node_9.add_father_node(node_4)
        node_4.add_next_node(node_9)
        node_9.add_next_node(node_7)
        node_7.add_father_node(node_9)

        test_graph.add_node(node_5)
        test_graph.add_node(node_1)
        test_graph.add_node(node_2)
        test_graph.add_node(node_3)
        test_graph.add_node(node_4)
        test_graph.add_node(node_6)
        test_graph.add_node(node_8)
        test_graph.add_node(node_9)
        test_graph.add_node(node_7)
        return test_graph

    @staticmethod
    def get_bn_graph_no_center():
        test_graph = Graph()
        node_1 = Node("add", input_num=1, out_shape=[200])
        node_1.atomic_type = "add"
        node_2 = Node("sqrt", input_num=1, out_shape=[200])
        node_2.atomic_type = "sqrt"
        node_3 = Node("divide", input_num=1, out_shape=[200])
        node_3.atomic_type = "divide"
        node_4 = Node("multiply", input_num=1, out_shape=[200])
        node_4.atomic_type = "multiply"
        node_5 = Node("multiply-2", input_num=1, out_shape=[200])
        node_5.atomic_type = "multiply"
        node_6 = Node("add-2", input_num=1, out_shape=[200])
        node_6.atomic_type = "add"
        # node_7 = Node("add-3", input_num=1, out_shape=[200])
        # node_7.atomic_type = "add"

        node_8 = Node("neg", input_num=1, out_shape=[200])
        node_8.atomic_type = "negative"
        node_9 = Node("multiply-3", input_num=1, out_shape=[200])
        node_9.atomic_type = "multiply"

        node_1.add_next_node(node_2)
        node_2.add_father_node(node_1)
        node_2.add_next_node(node_3)
        node_3.add_father_node(node_2)
        node_3.add_next_node(node_4)
        node_4.add_father_node(node_3)
        node_4.add_next_node(node_5)
        node_5.add_father_node(node_4)
        node_5.add_next_node(node_6)
        node_6.add_father_node(node_5)
        # node_6.add_next_node(node_7)
        # node_7.add_father_node(node_6)

        node_5.add_father_node(node_8)
        node_8.add_next_node(node_5)
        node_9.add_father_node(node_4)
        node_4.add_next_node(node_9)
        node_9.add_next_node(node_6)
        node_6.add_father_node(node_9)

        test_graph.add_node(node_5)
        test_graph.add_node(node_1)
        test_graph.add_node(node_2)
        test_graph.add_node(node_3)
        test_graph.add_node(node_4)
        # test_graph.add_node(node_7)
        test_graph.add_node(node_8)
        test_graph.add_node(node_9)
        test_graph.add_node(node_6)
        return test_graph

    def get_bn_graph_no_scale():
        test_graph = Graph()
        node_1 = Node("add", input_num=1, out_shape=[200])
        node_1.atomic_type = "add"
        node_2 = Node("sqrt", input_num=1, out_shape=[200])
        node_2.atomic_type = "sqrt"
        node_3 = Node("divide", input_num=1, out_shape=[200])
        node_3.atomic_type = "divide"
        # node_4 = Node("multiply", input_num=1, out_shape=[200])
        # node_4.atomic_type = "multiply"
        node_5 = Node("multiply-2", input_num=1, out_shape=[200])
        node_5.atomic_type = "multiply"
        node_6 = Node("add-2", input_num=1, out_shape=[200])
        node_6.atomic_type = "add"
        node_7 = Node("add-3", input_num=1, out_shape=[200])
        node_7.atomic_type = "add"

        node_8 = Node("neg", input_num=1, out_shape=[200])
        node_8.atomic_type = "negative"
        node_9 = Node("multiply-3", input_num=1, out_shape=[200])
        node_9.atomic_type = "multiply"

        node_1.add_next_node(node_2)
        node_2.add_father_node(node_1)
        node_2.add_next_node(node_3)
        node_3.add_father_node(node_2)
        node_3.add_next_node(node_5)
        node_5.add_father_node(node_3)
        # node_4.add_next_node(node_5)
        # node_5.add_father_node(node_4)
        node_5.add_next_node(node_6)
        node_6.add_father_node(node_5)
        node_6.add_next_node(node_7)
        node_7.add_father_node(node_6)

        node_5.add_father_node(node_8)
        node_8.add_next_node(node_5)
        node_9.add_father_node(node_3)
        node_3.add_next_node(node_9)
        node_9.add_next_node(node_7)
        node_7.add_father_node(node_9)

        test_graph.add_node(node_5)
        test_graph.add_node(node_1)
        test_graph.add_node(node_2)
        test_graph.add_node(node_3)
        # test_graph.add_node(node_4)
        test_graph.add_node(node_6)
        test_graph.add_node(node_8)
        test_graph.add_node(node_9)
        test_graph.add_node(node_7)

        return test_graph

    def get_bn_graph_no_center_no_scale():
        test_graph = Graph()
        node_1 = Node("add", input_num=1, out_shape=[200])
        node_1.atomic_type = "add"
        node_2 = Node("sqrt", input_num=1, out_shape=[200])
        node_2.atomic_type = "sqrt"
        node_3 = Node("divide", input_num=1, out_shape=[200])
        node_3.atomic_type = "divide"
        # node_4 = Node("multiply", input_num=1, out_shape=[200])
        # node_4.atomic_type = "multiply"
        node_5 = Node("multiply-2", input_num=1, out_shape=[200])
        node_5.atomic_type = "multiply"
        node_6 = Node("add-2", input_num=1, out_shape=[200])
        node_6.atomic_type = "add"
        # node_7 = Node("add-3", input_num=1, out_shape=[200])
        # node_7.atomic_type = "add"

        node_8 = Node("neg", input_num=1, out_shape=[200])
        node_8.atomic_type = "negative"
        node_9 = Node("multiply-3", input_num=1, out_shape=[200])
        node_9.atomic_type = "multiply"

        node_1.add_next_node(node_2)
        node_2.add_father_node(node_1)
        node_2.add_next_node(node_3)
        node_3.add_father_node(node_2)
        node_3.add_next_node(node_5)
        node_5.add_father_node(node_3)
        # node_4.add_next_node(node_5)
        # node_5.add_father_node(node_4)
        node_5.add_next_node(node_6)
        node_6.add_father_node(node_5)
        # node_6.add_next_node(node_7)
        # node_7.add_father_node(node_6)

        node_5.add_father_node(node_8)
        node_8.add_next_node(node_5)
        node_9.add_father_node(node_3)
        node_3.add_next_node(node_9)
        node_9.add_next_node(node_6)
        node_6.add_father_node(node_9)

        test_graph.add_node(node_5)
        test_graph.add_node(node_1)
        test_graph.add_node(node_2)
        test_graph.add_node(node_3)
        # test_graph.add_node(node_4)
        # test_graph.add_node(node_7)
        test_graph.add_node(node_8)
        test_graph.add_node(node_9)
        test_graph.add_node(node_6)

        return test_graph
